import datetime
import getopt
import sys

import numpy as np
import pandas as pd
from dateutil.relativedelta import relativedelta

from air_web.config.config import config
from air_web.data_platform import mysql_con, sql_engine
from air_web.dw.dws_common import AggCommon
from air_web.web_flask.dal.base_dal import EsBaseDal

con, cursor = mysql_con.get_cursor()


def commit_sql(case_sql, table_name):
    sql = """update {table_name} set data_rank = (case {case_sql} else data_rank end)""".format(
        table_name=table_name, case_sql=case_sql
    )
    cursor.execute(sql)
    con.commit()


def main(
    power_idx="aclr_res_power_96*",
    start_date=None,
    end_date=None,
    sort_field="p_kt",
):
    es_dal = EsBaseDal(config["ES_HOST"])
    IS_DAY_MAX_LIST = [1, 3] if config["is_total_max"] else [2, 3]
    rules = [
        ("data_time", "query", ">=", start_date),
        ("data_time", "query", "<=", end_date),
        ("is_day_max", "query", "in", IS_DAY_MAX_LIST),
    ]

    is_cal = "1" if config["c_cons_filter_is_cal"] else "*"
    power_idx = (
        config["POWERLOAD"]
        .replace("on5", "*")
        .replace("on7", "*")
        .replace("cal01", is_cal)
    )

    df = es_dal.query_dataframe(
        rules,
        power_idx,
        ts=start_date,
        te=end_date,
        doc_time_field="data_time",
    )

    # c_cons
    df = df.drop_duplicates("cons_no")
    df = df.sort_values(sort_field, ascending=False)
    cons_df = df
    cons_df["data_rank"] = cons_df.reset_index().index + 1
    cons_df = cons_df[["cons_no", "data_rank"]]
    c_cons = config.get("C_CONS", "c_cons")

    # 先like表结构
    sql_c = f"create table if not exists {c_cons}_zsot_temp_sort like {c_cons}"
    cursor.execute(sql_c)
    # 造好数据
    sql_s = f"select * from {c_cons}"
    res_df = pd.DataFrame(mysql_con.get(sql_s))
    merged_df = res_df.merge(
        cons_df, on="cons_no", how="left", suffixes=("_merged", "_df")
    )
    merged_df["data_rank"] = merged_df["data_rank_df"]
    merged_df["data_rank"] = merged_df["data_rank"].replace({np.nan: 99999999})
    merged_df.drop(["data_rank_merged", "data_rank_df"], axis=1, inplace=True)
    sql_t = f"truncate {c_cons}_zsot_temp_sort"
    cursor.execute(sql_t)
    merged_df = merged_df.replace({np.nan: None})
    sql_engine.update_df_by_id(merged_df, f"{c_cons}_zsot_temp_sort")
    # 然后重命名两张表
    sql_cb = f"alter table {c_cons} rename to {c_cons}_zsot_temp_sort_bak"
    cursor.execute(sql_cb)
    sql_sc = f"alter table {c_cons}_zsot_temp_sort rename to {c_cons}"
    cursor.execute(sql_sc)
    sql_db = f"drop table {c_cons}_zsot_temp_sort_bak"
    cursor.execute(sql_db)
    con.commit()

    # area_type_map
    type_df = df.groupby(["type_id"]).agg({sort_field: "sum"}).reset_index()
    type_df = type_df.sort_values(sort_field, ascending=False)
    type_id_list = type_df["type_id"].tolist()

    area_type_map = config.get("AREA_TYPE_MAP", "area_type_map")
    type_case_sql = ""
    for idx, type_id in enumerate(type_id_list):
        type_case_sql = "{type_case_sql}when type_id = '{type_id}' then {data_rank} ".format(
            type_case_sql=type_case_sql, type_id=type_id, data_rank=idx + 1
        )
        if idx % 5000 == 0:
            commit_sql(type_case_sql, area_type_map)
            type_case_sql = ""
    commit_sql(type_case_sql, area_type_map)

    cursor.close()
    con.close()


class ConsSort(AggCommon):
    def __init__(
        self,
        logger,
        task,
        start_time,
        end_time,
        step,
        step_dict,
        timedelta,
        rewrite,
    ):
        super().__init__(
            logger,
            task,
            start_time,
            end_time,
            step,
            step_dict,
            timedelta,
            rewrite,
        )

    def main(
        self,
        power_idx="aclr_res_power_96*",
        sort_field="p_kt",
        start_date=None,
        end_date=None,
    ):
        if end_date is not None:
            self.end_time = datetime.datetime.strptime(
                end_date, "%Y-%m-%d %H:%M:%S"
            )
        if start_date is not None:
            self.start_time = datetime.datetime.strptime(
                start_date, "%Y-%m-%d %H:%M:%S"
            )
        else:
            self.start_time = self.end_time - relativedelta(days=7)

        start_date = self.start_time
        end_date = self.end_time
        main(power_idx, start_date, end_date, sort_field)


if __name__ == "__main__":
    power_idx = "aclr_res_power_96*"
    end_date = datetime.datetime.now().strftime("%Y-%m-%d")
    start_date = (end_date - relativedelta(days=7)).strftime("%Y-%m-%d")
    sort_field = "p_kt"

    opts, args = getopt.getopt(sys.argv[1:], "s:e:c:")
    for opt, val in opts:
        if opt == "-s":
            start_date = val
        elif opt == "-e":
            end_date = val
        elif opt == "-c":
            sort_field = val

    main(power_idx, start_date, end_date, sort_field)
